/**
 * @param nprobe            Faiss(IVFFLAT)构建参数：倒排列表（Inverted List）数量
 * @param metric            Faiss(IVFFLAT)构建参数：距离度量类型
 * @param quantizer         Faiss(IVFFLAT)量化器
 * @param index             Faiss(IVFFLAT)索引
 * @param ivf_params        Faiss(IVFFLAT)ivf参数：用于添加nprobe
 */

#ifndef FAISS_IVFFLAT_INDEX_H
#define FAISS_IVFFLAT_INDEX_H

#include "index_base.h"
#include "datasets.h"
#include "config_parser.h"
#include <faiss/IndexFlat.h>
#include <faiss/IndexIVFFlat.h>
#include <faiss/IndexRefine.h>
#include <memory>

class FaissIVFFLATIndex : public IndexBase {
public:
    FaissIVFFLATIndex(const ConfigParser &config);
    ~FaissIVFFLATIndex();
    virtual void build(const Datasets *data) override;
    virtual void prepare_search() override;
    virtual void search_single(const float *xq, int k, float *distances, int64_t *labels) override;
    virtual void search_batch(const int nq, const float *xq, int k, float *distances, int64_t *labels) override;

private:
    int nprobe;
    faiss::MetricType metric;
    std::unique_ptr<faiss::IndexFlat> quantizer;
    std::unique_ptr<faiss::IndexIVFFlat> index;
    faiss::SearchParametersIVF ivf_params;
    std::string index_path;
    std::string save_or_load;
};

#endif  // FAISS_IVFFLAT_INDEX_H
